/* introspect the CPU host. receive VM's GPA and send back the corresponding data */
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/mman.h>
#include <stdio.h>
#include <inttypes.h>
#include <getopt.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <pthread.h>

#include <libvmi/libvmi.h>

#define HOST_NAME               "host"
#define DEFAULT_LOG_FILENAME    "daemon.log"

#define SOCKET_PATH             "/tmp/vmi.sock"

struct kvm_struct_offest {
    uint32_t linux_mm_pgd;
    uint32_t linux_kvm;
    uint32_t linux_kvm_pid;
    uint32_t linux_kvm_mm;
    uint32_t linux_kvm_memslots0;
    uint32_t linux_kvm_slots0_gfntree;
    uint32_t linux_kvm_slots0_nodeidx;
    uint32_t linux_kvm_slot_gfnode0;
    uint32_t linux_kvm_slot_gfnode1;
    uint32_t linux_kvm_slot_basegfn;
    uint32_t linux_kvm_slot_nrpages;
    uint32_t linux_kvm_slot_hva;
    uint32_t linux_kvm_rbnode_right;
    uint32_t linux_kvm_rbnode_left;
};

struct kvm_struct_offest g_kvm_offset = { /* can be embedded into /etc/libvmi.conf for Host */
    .linux_mm_pgd = 0x80,
    .linux_kvm = 0x9a0,
    .linux_kvm_pid = 0xc40,
    .linux_kvm_mm = 0x98,
    .linux_kvm_memslots0 = 0x908,
    .linux_kvm_slots0_gfntree = 0x20,
    .linux_kvm_slots0_nodeidx = 0x428,
    .linux_kvm_slot_gfnode0 = 0x80,
    .linux_kvm_slot_gfnode1 = 0x98,
    .linux_kvm_slot_basegfn = 0xb0,
    .linux_kvm_slot_nrpages = 0xb8,
    .linux_kvm_slot_hva = 0xc8,
    .linux_kvm_rbnode_right = 0x8,
    .linux_kvm_rbnode_left = 0x10
};

#define KVM_MEMSLOT_MAX_NUM     10
struct kvm_slot {
    addr_t gpa_start;
    addr_t gpa_end;
    addr_t hva;
};

struct kvm_vm {
    addr_t st_kvm;
    vmi_pid_t pid;
    addr_t user_pgd;
    addr_t gpa_start;
    addr_t gpa_end;
    uint32_t slot_num;
    struct kvm_slot slot[KVM_MEMSLOT_MAX_NUM];
};

#define VMI_MESSAGE_MAX_ARGS    6
struct vmi_message {
    int msg_type;
    int msg_argc;
    uint64_t msg_args[VMI_MESSAGE_MAX_ARGS];
};

enum vmi_message_type {
    VMI_MSG_GET_MEM_RANGE = 0,
    VMI_MSG_GET_GPA_DATA,
    VMI_MSG_ERR,
};

int parse_args(int argc, char **argv, uint64_t *vmid, uint64_t *gpa, uint64_t *size, FILE **log_fp) {
    char *output_filename = NULL;

    if ( argc < 2 ) {
        printf("Usage: %s\n", argv[0]);
        printf("\t -o/--output-file <log file>\n");
        printf("\t -g/--gpa <GPA>\n");
        printf("\t -s/--size <size of data to be read>\n");
        printf("\t -v/--vm-id <VM id>\n");
        return 0;
    }

    if ( argc > 2 ) {
        const struct option long_opts[] = {
            {"vm-id", required_argument, NULL, 'v'},
            {"gpa", optional_argument, NULL, 'g'},
            {"size", optional_argument, NULL, 's'},
            {"output-file", optional_argument, NULL, 'o'},
            {NULL, 0, NULL, 0}
        };
        const char* opts = "v:g:o:s:";
        int c;
        int long_index = 0;

        while ((c = getopt_long (argc, argv, opts, long_opts, &long_index)) != -1) {
            switch (c) {
                case 'v':
                    *vmid = strtoull(optarg, NULL, 0);
                    break;
                case 'g':
                    *gpa = strtoull(optarg, NULL, 0);
                    break;
                case 's':
                    *size = strtoull(optarg, NULL, 0);
                    break;
                case 'o':
                    output_filename = optarg;
                    break;
                default:
                    printf("Unknown option\n");
                    return -1;
            }
        }
    }

    if (output_filename != NULL) {
        *log_fp = fopen(output_filename, "w+");
    } else {
        *log_fp = fopen(DEFAULT_LOG_FILENAME, "w+");
    }
    if (*log_fp == NULL) {
        printf("cannot write to output file: %s\n", output_filename ? output_filename : DEFAULT_LOG_FILENAME);
        return -1;
    }
    return 0;
}

int test_process_list(vmi_instance_t vmi, FILE *log_fp) {
    vmi_mode_t mode;
    addr_t current_process = 0;
    char *procname = NULL;
    vmi_pid_t pid = 0;
    unsigned long tasks_offset = 0, pid_offset = 0, name_offset = 0;
    addr_t list_head = 0, cur_list_entry = 0, next_list_entry = 0;

    fprintf(log_fp, ">>> luowu. os is linux. try to vmi_get_offset\n");
    if ( VMI_FAILURE == vmi_get_offset(vmi, "linux_tasks", &tasks_offset) )
        return -1;
    if ( VMI_FAILURE == vmi_get_offset(vmi, "linux_name", &name_offset) )
        return -1;
    if ( VMI_FAILURE == vmi_get_offset(vmi, "linux_pid", &pid_offset) )
        return -1;

    /* demonstrate name and id accessors */
    if (VMI_FAILURE == vmi_get_access_mode(vmi, NULL, 0, NULL, &mode))
        return -1;

    if ( VMI_FILE != mode ) {
        uint64_t id = vmi_get_vmid(vmi);
        fprintf(log_fp, "Process listing for Host %s (id=%"PRIu64")\n", HOST_NAME, id);
    } else {
        fprintf(log_fp, "Process listing for file %s\n", HOST_NAME);
    }
    fflush(log_fp);

    /* get the head of the list */
    if ( VMI_FAILURE == vmi_translate_ksym2v(vmi, "init_task", &list_head) )
        return -1;

    list_head += tasks_offset;
    cur_list_entry = list_head;
    if (VMI_FAILURE == vmi_read_addr_va(vmi, cur_list_entry, 0, &next_list_entry)) {
        fprintf(log_fp, "Failed to read next pointer at %"PRIx64"\n", cur_list_entry);
        return -1;
    }

    /* walk the task list */
    while (1) {
        current_process = cur_list_entry - tasks_offset;
        vmi_read_32_va(vmi, current_process + pid_offset, 0, (uint32_t*)&pid);

        procname = vmi_read_str_va(vmi, current_process + name_offset, 0);
        if (!procname) {
            fprintf(log_fp, "Failed to find procname\n");
            return -1;
        }

        /* print out the process name */
        fprintf(log_fp, "[%5d] %s (struct addr:%"PRIx64")\n", pid, procname, current_process);
        fflush(log_fp);
        if (procname) {
            free(procname);
            procname = NULL;
        }

        /* follow the next pointer */
        cur_list_entry = next_list_entry;
        if (VMI_FAILURE == vmi_read_addr_va(vmi, cur_list_entry, 0, &next_list_entry)) {
            fprintf(log_fp, "Failed to read next pointer in loop at %"PRIx64"\n", cur_list_entry);
            return -1;
        }

        if (cur_list_entry == list_head)
            break;
    };
    return 0;
}

#define VM_SLOT_MAX_NUMBER  100
int get_vm_gpa_range(vmi_instance_t vmi, FILE *log_fp, struct kvm_vm *found_vm) {
    /* reference KVM gfn_to_hva */
    addr_t node_queue[VM_SLOT_MAX_NUMBER] = { 0 };
    uint32_t node_idx = 0, node_len = 0;
    addr_t cur_list_entry = 0, next_list_entry = 0;
    addr_t slots = 0;
    int idx = 0;

    if (found_vm == 0) {
        fprintf(log_fp, "[%s] found_vm not initialized.\n", __func__);
        fflush(log_fp);
        return -1;
    }

    /* slots = kvm->memslots[0] */
    if (VMI_FAILURE == vmi_read_addr_va(vmi, found_vm->st_kvm + g_kvm_offset.linux_kvm_memslots0, 0, &slots)) {
        fprintf(log_fp, "[%s] Failed to get memslots at %"PRIx64"\n", __func__, found_vm->st_kvm + g_kvm_offset.linux_kvm_memslots0);
        fflush(log_fp);
        return -1;
    }

    /* idx = slots->node_idx; */
    vmi_read_32_va(vmi, slots + g_kvm_offset.linux_kvm_slots0_nodeidx, 0, (uint32_t*)&idx);
    if (idx != 0 && idx != 1) {
        fprintf(log_fp, "[%s] bad idx[%d] in memslots %"PRIx64"\n", __func__, idx, slots + g_kvm_offset.linux_kvm_slots0_nodeidx);
        fflush(log_fp);
        return -1;
    }

    /*  for (node = slots->gfn_tree.rb_node; node; ) */
    if (VMI_FAILURE == vmi_read_addr_va(vmi, slots + g_kvm_offset.linux_kvm_slots0_gfntree, 0, &(node_queue[node_idx]))) {
        fprintf(log_fp, "[%s] Failed to get slots->gfn_tree.rb_node at %"PRIx64"\n", __func__, slots + g_kvm_offset.linux_kvm_slots0_gfntree);
        fflush(log_fp);
        return -1;
    }
    node_len += 1;

    while (node_idx < node_len) {
        addr_t base_gfn = 0, gpa_start = 0, gpa_end = 0, hva = 0;
        uint64_t nr_pages = 0;

        cur_list_entry = node_queue[node_idx++];
        if (cur_list_entry == 0)
            break;

        /* slot = container_of(node, struct kvm_memory_slot, gfn_node[idx]); */
        addr_t slot = idx == 0 ? cur_list_entry - g_kvm_offset.linux_kvm_slot_gfnode0 :
            cur_list_entry - g_kvm_offset.linux_kvm_slot_gfnode1;

        vmi_read_64_va(vmi, slot + g_kvm_offset.linux_kvm_slot_basegfn, 0, (uint64_t*)&base_gfn);
        vmi_read_64_va(vmi, slot + g_kvm_offset.linux_kvm_slot_nrpages, 0, (uint64_t*)&nr_pages);
        vmi_read_64_va(vmi, slot + g_kvm_offset.linux_kvm_slot_hva, 0, (uint64_t*)&hva);

        /* update kvm VM info */
        gpa_start = (uint64_t)base_gfn << 12;
        gpa_end = ((uint64_t)base_gfn + nr_pages) << 12;
        if (found_vm->gpa_start == (addr_t)-1) {
            found_vm->gpa_start = gpa_start;
            found_vm->gpa_end = gpa_end;
        } else {
            found_vm->gpa_start = found_vm->gpa_start > gpa_start ? gpa_start : found_vm->gpa_start;
            found_vm->gpa_end = found_vm->gpa_end < gpa_end ? gpa_end : found_vm->gpa_end;
        }
        if (found_vm->slot_num >= KVM_MEMSLOT_MAX_NUM) {
            fprintf(log_fp, "[%s] Found too many slots [%d] for VM\n", __func__, found_vm->slot_num);
            fflush(log_fp);
            return -1;
        }
        found_vm->slot[found_vm->slot_num].gpa_start = gpa_start;
        found_vm->slot[found_vm->slot_num].gpa_end = gpa_end;
        found_vm->slot[found_vm->slot_num].hva = hva;
        found_vm->slot_num++;

        fprintf(log_fp, "[%s] found GFN:0x%lx nr_page: 0x%x hva: 0x%lx. VM's gpa: 0x%lx-0x%lx\n", __func__,
            base_gfn, nr_pages, hva, found_vm->gpa_start, found_vm->gpa_end);
        fflush(log_fp);

        /* node = node->rb_right */
        if (VMI_FAILURE == vmi_read_addr_va(vmi, cur_list_entry + g_kvm_offset.linux_kvm_rbnode_right, 0, &next_list_entry)) {
            fprintf(log_fp, "[%s] Failed to get node->rb_right at %"PRIx64"\n", __func__, cur_list_entry + g_kvm_offset.linux_kvm_rbnode_right);
            fflush(log_fp);
            return -1;
        }
        if (next_list_entry != 0) {
            if (node_len >= VM_SLOT_MAX_NUMBER) {
                fprintf(log_fp, "[%s] found too many nodes. You should increase the length of node_queue\n", __func__);
                fflush(log_fp);
                return -1;
            }
            node_queue[node_len++] = next_list_entry;
        }

        /* node = node->rb_left */
        if (VMI_FAILURE == vmi_read_addr_va(vmi, cur_list_entry + g_kvm_offset.linux_kvm_rbnode_left, 0, &next_list_entry)) {
            fprintf(log_fp, "[%s] Failed to get node->rb_left at %"PRIx64"\n", __func__, cur_list_entry + g_kvm_offset.linux_kvm_rbnode_left);
            fflush(log_fp);
            return -1;
        }
        if (next_list_entry != 0) {
            if (node_len >= VM_SLOT_MAX_NUMBER) {
                fprintf(log_fp, "[%s] found too many nodes. You should increase the length of node_queue\n", __func__);
                fflush(log_fp);
                return -1;
            }
            node_queue[node_len++] = next_list_entry;
        }
        
        return 0;
    }
}

int inspect_vm_list(vmi_instance_t vmi, FILE *log_fp, struct kvm_vm *found_vm) {
    addr_t list_head = 0, cur_list_entry = 0, next_list_entry = 0;
    unsigned long kvm2vmlist_offset = g_kvm_offset.linux_kvm, kvm2pid_offset = g_kvm_offset.linux_kvm_pid;
    vmi_pid_t pid = 0;
    addr_t current_kvm = 0, mm = 0, pgd_ptr = 0, pgd = 0;
    bool find_vm = false;

    /* get the head of the list */
    if ( VMI_FAILURE == vmi_translate_ksym2v(vmi, "vm_list", &list_head) )
        return -1;

    fprintf(log_fp, ">>> luowu. Found vm_list in 0x%lx. kvm2vmlist_offset: 0x%lx, kvm2pid_offset: 0x%lx\n",
        list_head, kvm2vmlist_offset, kvm2pid_offset);
    fflush(log_fp);

    cur_list_entry = list_head;
    while (1) {
        if (VMI_FAILURE == vmi_read_addr_va(vmi, cur_list_entry, 0, &next_list_entry)) {
            fprintf(log_fp, "Failed to read next pointer at %"PRIx64"\n", cur_list_entry);
            fflush(log_fp);
            return -1;
        }

        if (next_list_entry == list_head) {
            if (find_vm) {
                break;
            }
            fprintf(log_fp, "No VM found in KVM's vm-list. Sleep 10s...\n");
            fflush(log_fp);
            sleep(10);
            continue;
        }

        /* kvm->vm_list --- kvm */
        current_kvm = next_list_entry - kvm2vmlist_offset;
        /* kvm --- kvm_pid */
        vmi_read_32_va(vmi, current_kvm + kvm2pid_offset, 0, (uint32_t*)&pid);
        /* kvm --- mm */
        if (VMI_FAILURE == vmi_read_addr_va(vmi, current_kvm + g_kvm_offset.linux_kvm_mm, 0, &mm)) {
            fprintf(log_fp, "Failed to read kvm --- mm at %"PRIx64"\n", current_kvm + g_kvm_offset.linux_kvm_mm);
            fflush(log_fp);
            return -1;
        }
        /* mm --- pgd */
        if (VMI_FAILURE == vmi_read_addr_va(vmi, mm + g_kvm_offset.linux_mm_pgd, 0, &pgd_ptr)) {
            fprintf(log_fp, "Failed to read mm --- pgd at %"PRIx64"\n", mm + g_kvm_offset.linux_mm_pgd);
            fflush(log_fp);
            return -1;
        }
        fprintf(log_fp, "Found VM: pid[%5d](struct addr:%"PRIx64", pgd_ptr:0x%lx)\n", pid, current_kvm, pgd_ptr);
        fflush(log_fp);

        found_vm->st_kvm = current_kvm;
        found_vm->pid = pid;
        found_vm->user_pgd = pgd_ptr;

        find_vm = true;
        /* follow the next pointer */
        cur_list_entry = next_list_entry;
    }

    if (found_vm)
        return get_vm_gpa_range(vmi, log_fp, found_vm);
    return -1;
}

struct thread_args {
    vmi_instance_t vmi;
    FILE *log_fp;
    int client_sock;
    struct kvm_vm *found_vm;
};

#define GET_VA_VPN4(x)                      ((x & ((1 << 57) - 1)) >> 48)
#define GET_VA_VPN3(x)                      ((x & ((1 << 48) - 1)) >> 39)
#define GET_VA_VPN2(x)                      ((x & ((1 << 39) - 1)) >> 30)
#define GET_VA_VPN1(x)                      ((x & ((1 << 30) - 1)) >> 21)
#define GET_VA_VPN0(x)                      ((x & ((1 << 21) - 1)) >> 12)
#define GET_VA_OFFSET(x)                    (x & ((1 << 12) - 1))

#define GET_PAGE_SIZE_IN_LEVEL(level)       (1 << (12 + 9 * level))
#define GET_VA_OFFSET_IN_LEVEL(va, level)   (va & ((1 << (12 + 9 * level)) - 1))

#define GET_VA_VPN_IN_LEVEL(va, level)      ((va & (((addr_t)1 << (21 + 9 * level)) - 1)) >> (12 + 9 * level))

#define GET_PPN_FROM_PTE(x)                 (((x >> 10) & (((addr_t)1 << 44) - 1)) << 12)
#define GET_V_BIT_FROM_PTE(x)               (x & 0x1)
#define GET_R_BIT_FROM_PTE(x)               (x & 0x2)
#define GET_W_BIT_FROM_PTE(x)               (x & 0x4)
#define GET_X_BIT_FROM_PTE(x)               (x & 0x8)

#define PTE_SIZE                8

static int get_pa_by_page_walk(struct thread_args *args, addr_t vaddr, addr_t *found_pa) {
    struct kvm_vm *vm = args->found_vm;
    int fd = args->client_sock;
    int level = 5; // SV57
    int i = 0;
    addr_t pd = vm->user_pgd;

    fprintf(args->log_fp, ">>> luowu. client fd: %d. Try to page walk for pid [%d]. pgd=0x%lx hva=0x%lx\n", fd, vm->pid, vm->user_pgd, vaddr);
    fflush(args->log_fp);

    fprintf(args->log_fp, ">>> luowu. client fd: %d. --RISCV64 SV57 PTLookup: vaddr = 0x%.16"PRIx64", pt = 0x%.16"PRIx64"\n", fd, vaddr, pd);
    while(level-- > 0) {
        addr_t vpn = GET_VA_VPN_IN_LEVEL(vaddr, level);
        addr_t pte = pd + vpn * PTE_SIZE;
        uint64_t pte_data = 0;
        uint64_t ppn = 0;

        fprintf(args->log_fp, ">>> luowu. client fd: %d. --search for level %d. pd: 0x%lx. vpn: 0x%lx. pte_addr: 0x%lx\n", fd, level, pd, vpn, pte);

        // for pgd, we use vmi_read_64_va else we use vmi_read_64_pa
        if (level == 4) {
            if (VMI_FAILURE == vmi_read_64_va(args->vmi, 0, pte, &pte_data)) {
                fprintf(args->log_fp, ">>> luowu. client fd: %d. cannot read VA: 0x%lx\n", fd, pte);
                return -1;
            }
        } else {
            if (VMI_FAILURE == vmi_read_64_pa(args->vmi, pte, &pte_data)) {
                fprintf(args->log_fp, ">>> luowu. client fd: %d. cannot read PA: 0x%lx\n", fd, pte);
                return -1;
            }
        }

        ppn = GET_PPN_FROM_PTE(pte_data);
        if (GET_V_BIT_FROM_PTE(pte_data) != 1) {
            fprintf(args->log_fp, ">>> luowu. client fd: %d. bad pte data: 0x%lx in pa 0x%lx\n", fd, pte_data, pte);
            return -1;
        }

        if (level == 0) {
            // the last layer
            *found_pa = ppn + GET_VA_OFFSET(vaddr);
            fprintf(args->log_fp, ">>> luowu. client fd: %d. --Found PA in level %d. pa: 0x%lx\n", fd, level, *found_pa);
            return 0;
        }

        fprintf(args->log_fp, ">>> luowu. client fd: %d. -- ptedata=0x%lx ppn=0x%lx\n", fd, pte_data, ppn);

        if (GET_R_BIT_FROM_PTE(pte_data) != 0 || GET_X_BIT_FROM_PTE(pte_data) != 0) {
            // found the leaf node
            *found_pa = ppn + GET_VA_OFFSET_IN_LEVEL(vaddr, level);
            fprintf(args->log_fp, ">>> luowu. client fd: %d. --Found PA in level %d. pa: 0x%lx\n", fd, level, *found_pa);
            return 0;
        }

        pd = ppn;
    }
    fprintf(args->log_fp, ">>> luowu. client fd: %d. --Not Found a leaf PT\n", fd);

    return -1;
}

static int handle_get_gpa_msg(struct thread_args *args, struct vmi_message *msg) {
    addr_t gpa = msg->msg_args[0];
    uint32_t length = msg->msg_args[1];
    struct kvm_vm *vm = args->found_vm;
    int fd = args->client_sock;
    int i;
    addr_t hva = 0;

    if (vm->slot_num <= 0) {
        fprintf(args->log_fp, ">>> luowu. client fd: %d. Error: does not found any slot\n", fd);
        fflush(args->log_fp);
        goto err;
    }

    /* which slot the GPA belong to, and determine HVA */
    for (i = 0; i < vm->slot_num; i++) {
        if (gpa >= vm->slot[i].gpa_start && gpa <= vm->slot[i].gpa_end) {
            hva = vm->slot[i].hva + (gpa - vm->slot[i].gpa_start);
            fprintf(args->log_fp, ">>> luowu. client fd: %d. GPA 0x%lx is in Slot %d: hva: 0x%lx gpa: 0x%lx-0x%lx\n", fd,
                gpa, i, vm->slot[i].hva, vm->slot[i].gpa_start, vm->slot[i].gpa_end);
            fflush(args->log_fp);
            break;
        }
    }

    if (hva == 0) {
        fprintf(args->log_fp, ">>> luowu. client fd: %d. Error: does not found gpa 0x%lx in all slots\n", fd, gpa);
        fflush(args->log_fp);
        goto err;
    }

    uint32_t out_len = 0;
    char *buf = malloc(sizeof(struct vmi_message) + length);
    memset(buf, 0, sizeof(struct vmi_message) + length);
    fprintf(args->log_fp, ">>> luowu. client fd: %d. try to vmi_read_va: 0x%lx len:0x%x\n", fd, hva, length);
    fflush(args->log_fp);
    if (VMI_FAILURE == vmi_read_va(args->vmi, hva, vm->pid, length, buf + sizeof(struct vmi_message), &out_len) ||
        out_len != length) {
        fprintf(args->log_fp, ">>> luowu. client fd: %d. cannot vmi_read_va: 0x%lx len:0x%x\n", fd, hva, length);
        goto err;
    }

    // /* page walk the userspace process to get PA */
    // addr_t found_pa = 0;
    // if (get_pa_by_page_walk(args, hva, &found_pa) != 0) {
    //     fprintf(args->log_fp, ">>> luowu. client fd: %d. Error: cannot page walk for HVA: 0x%lx\n", fd, hva);
    //     fflush(args->log_fp);
    //     return -1;
    // }

    // /* read the content */
    // fprintf(args->log_fp, ">>> luowu. client fd: %d. try to vmi_read_pa after page walk: 0x%lx len:0x%x\n", fd, found_pa, length);
    // fflush(args->log_fp);
    // if (VMI_FAILURE == vmi_read_pa(args->vmi, found_pa, length, buf + sizeof(struct vmi_message), &out_len) ||
    //     out_len != length) {
    //     fprintf(args->log_fp, ">>> luowu. client fd: %d. cannot vmi_read_pa: 0x%lx len:0x%x\n", fd, found_pa, length);
    //     return -1;
    // }

    fprintf(args->log_fp, ">>> luowu. client fd: %d. found content in GPA: 0x%lx. Reply it!\n", fd, gpa);
    fflush(args->log_fp);
    /* set message body and sent it back */
    struct vmi_message *reply = (struct vmi_message*)buf;
    reply->msg_type = VMI_MSG_GET_GPA_DATA;
    reply->msg_argc = 2;
    reply->msg_args[0] = gpa;
    reply->msg_args[1] = length;

    fprintf(args->log_fp, ">>> luowu. client fd: %d. [%s] for VM. data replied by VMI Daemon is: 0x%lx (show 8 bytes)\n", fd, __func__,
            length > 8 ? *(uint64_t*)buf : 0);
    fflush(args->log_fp);

    write(args->client_sock, buf, sizeof(struct vmi_message) + length);
    free(buf);
    return 0;
err:
    struct vmi_message reply_err;
    memset((char*)&reply_err, 0, sizeof(struct vmi_message));
    reply_err.msg_type = VMI_MSG_ERR;
    write(args->client_sock, (char*)&reply_err, sizeof(struct vmi_message));
    return -1;
}

int handle_msg(void *data) {
    struct thread_args *args = (struct thread_args*)data;
    char buffer[1024];

    fprintf(args->log_fp, ">>> luowu. Server start a new thread to handle msgs for client: %d\n", args->client_sock);
    fflush(args->log_fp);

    while (1) {
        ssize_t num_bytes = read(args->client_sock, buffer, sizeof(buffer));
        if (num_bytes == -1) {
            fprintf(args->log_fp, ">>> luowu. client fd: %d. Error: Read failed\n", args->client_sock);
            fflush(args->log_fp);
            return -1;
        }
        if (num_bytes != sizeof(struct vmi_message)) {
            fprintf(args->log_fp, ">>> luowu. client fd: %d. bad msg\n", args->client_sock);
            fflush(args->log_fp);
            return -1;
        }

        struct vmi_message msg, reply;
        memset(&reply, 0, sizeof(reply));
        memcpy(&msg, buffer, sizeof(msg));

        switch (msg.msg_type)
        {
        case VMI_MSG_GET_MEM_RANGE:
            fprintf(args->log_fp, ">>> luowu. client fd: %d. Received message to get mem range. Reply it!\n", args->client_sock);
            fflush(args->log_fp);
            reply.msg_type = VMI_MSG_GET_MEM_RANGE;
            reply.msg_argc = 2;
            reply.msg_args[0] = args->found_vm->gpa_start;
            reply.msg_args[1] = args->found_vm->gpa_end;
            write(args->client_sock, (char*)&reply, sizeof(reply));
            break;
        case VMI_MSG_GET_GPA_DATA:
            fprintf(args->log_fp, ">>> luowu. client fd: %d. Received message to get GPA data\n", args->client_sock);
            fflush(args->log_fp);
            (void)handle_get_gpa_msg(args, &msg);
            break;
        default:
            fprintf(args->log_fp, ">>> luowu. client fd: %d. Received Non-Implemented message. type:%d\n", args->client_sock, msg.msg_type);
            fflush(args->log_fp);
            close(args->client_sock);
            return -1;
        }
    }

    close(args->client_sock);
    return 0;
}

int start_listener(vmi_instance_t vmi, FILE *log_fp, struct kvm_vm *found_vm) {
    int server_sock = socket(AF_UNIX, SOCK_STREAM, 0);
    if (server_sock == -1) {
        fprintf(log_fp, ">>> luowu. Error: Socket creation failed");
        return -1;
    }
    struct sockaddr_un server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sun_family = AF_UNIX;
    strncpy(server_addr.sun_path, SOCKET_PATH, sizeof(server_addr.sun_path) - 1);

    if (bind(server_sock, (struct sockaddr*)&server_addr, sizeof(struct sockaddr_un)) == -1) {
        fprintf(log_fp, ">>> luowu. Error: Socket binding failed");
        close(server_sock);
        return -1;
    }

    if (listen(server_sock, 5) == -1) {
        fprintf(log_fp, ">>> luowu. Error: Listen failed");
        close(server_sock);
        return -1;
    }
    fprintf(log_fp, ">>> luowu. Server started and waiting for connections...\n");
    fflush(log_fp);
    while (1) {
        int client_sock = accept(server_sock, NULL, NULL);
        if (client_sock == -1) {
            fprintf(log_fp, ">>> luowu. Error: Accept failed");
            continue;
        }

        struct thread_args args = {
            .client_sock = client_sock,
            .found_vm = found_vm,
            .log_fp = log_fp,
            .vmi = vmi
        };
        pthread_t t;
        pthread_create(&t, NULL, handle_msg, &args);
    }
    close(server_sock);
    return 0;
}

int main (int argc, char **argv) {
    vmi_instance_t vmi = {0};
    uint64_t gpa = 0, size = 0, vmid = 0;
    FILE *log_fp = NULL;

    int rc = -1;
    rc = parse_args(argc, argv, &vmid, &gpa, &size, &log_fp);
    if (rc != 0) {
        fprintf(log_fp, "Failed to parse args.\n");
        return -1;
    }

    /* initialize the libvmi library */
    fprintf(log_fp, ">>> luowu. try to vmi_init_complete: find kpgd, kaslr, init_task...\n");
    fflush(log_fp);
    if (VMI_FAILURE == vmi_init_complete(&vmi, HOST_NAME, VMI_INIT_DOMAINNAME, NULL, VMI_CONFIG_GLOBAL_FILE_ENTRY, NULL, NULL)) {
        fprintf(log_fp, "Failed to init LibVMI library.\n");
        goto error_exit;
    }

    /* init the offset values */
    if (VMI_OS_LINUX != vmi_get_ostype(vmi)) {
        fprintf(log_fp, "Error: ostype is not Linux.\n");
        goto error_exit;
    }

    fprintf(log_fp, ">>> luowu. test... show the %s's process list...\n", HOST_NAME);
    fflush(log_fp);
    rc = test_process_list(vmi, log_fp);
    if (rc != 0) {
        fprintf(log_fp, "Failed to get process list.\n");
        return -1;
    }

    fprintf(log_fp, ">>> luowu. inspect VM list...\n");
    fflush(log_fp);
    struct kvm_vm found_vm;
    memset(&found_vm, 0, sizeof(struct kvm_vm));
    found_vm.gpa_start = (addr_t)-1;
    rc = inspect_vm_list(vmi, log_fp, &found_vm);
    if (rc != 0) {
        fprintf(log_fp, "Failed to inspect VM list.\n");
        return -1;
    }

    fprintf(log_fp, ">>> luowu. start listener for VMI agent...\n");
    fflush(log_fp);
    rc = start_listener(vmi, log_fp, &found_vm);

error_exit:
    /* cleanup any memory associated with the LibVMI instance */
    vmi_destroy(vmi);
    return rc;
}
